import torch
from torch import nn

class CA_attention(nn.Module):
    def __init__(self, channel, h, w, reduction=16):
        super(CA_attention, self).__init__()

        self.h = h
        self.w = w

		# 在水平和垂直方向进行 平均池化
        self.avg_pool_x = nn.AdaptiveAvgPool2d((h, 1))
        self.avg_pool_y = nn.AdaptiveAvgPool2d((1, w))

		# 提取通道特征
        self.conv_1x1 = nn.Conv2d(in_channels=channel, out_channels=channel // reduction, kernel_size=1, stride=1,
                                  bias=False)

        self.relu = nn.ReLU()
        self.bn = nn.BatchNorm2d(channel // reduction)

        self.F_h = nn.Conv2d(in_channels=channel // reduction, out_channels=channel, kernel_size=1, stride=1,
                             bias=False)
        self.F_w = nn.Conv2d(in_channels=channel // reduction, out_channels=channel, kernel_size=1, stride=1,
                             bias=False)

        self.sigmoid_h = nn.Sigmoid()
        self.sigmoid_w = nn.Sigmoid()

    def forward(self, x):
        x_h = self.avg_pool_x(x).permute(0, 1, 3, 2)
        x_w = self.avg_pool_y(x)
		# 加入标准化
        x_cat_conv_relu = self.relu(self.bn(self.conv_1x1(torch.cat((x_h, x_w), 3))))

        x_cat_conv_split_h, x_cat_conv_split_w = x_cat_conv_relu.split([self.h, self.w], 3)

        s_h = self.sigmoid_h(self.F_h(x_cat_conv_split_h.permute(0, 1, 3, 2)))
        s_w = self.sigmoid_w(self.F_w(x_cat_conv_split_w))

        out = x * s_h.expand_as(x) * s_w.expand_as(x)

        return out

# 通道域
class ChannelAttention(nn.Module):
    def __init__(self, in_channels, reduction_ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction_ratio),
            nn.ReLU(inplace=True),
            nn.Linear(in_channels // reduction_ratio, in_channels)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * self.sigmoid(y)

# 空间域
class SpatialAttention(nn.Module):
    def __init__(self):
        super(SpatialAttention, self).__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size=7, padding=3)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_pool = torch.max(x, dim=1, keepdim=True)[0]
        avg_pool = torch.mean(x, dim=1, keepdim=True)
        y = torch.cat([max_pool, avg_pool], dim=1)
        y = self.conv(y)
        return x * self.sigmoid(y)


class CBAM_attention(nn.Module):
    def __init__(self, in_channels, reduction=16):
        super(CBAM_attention, self).__init__()
        self.channel_attention = ChannelAttention(in_channels, reduction)
        self.spatial_attention = SpatialAttention()

    def forward(self, x):
        x = self.channel_attention(x)
        x = self.spatial_attention(x)
        return x


class SE_attention(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SE_attention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y


class CNN_MODEL(nn.Module):
	def __init__(self):
		super(CNN_MODEL, self).__init__()
		self.Frist_layer = nn.Sequential(
			nn.Conv2d(3, 32, 3, padding=1, stride=1, padding_mode='reflect'),
			nn.BatchNorm2d(32),
			nn.ReLU(),
			nn.Dropout(0.3),
			# 206 d
		)
		self.hidden_layer = nn.Sequential(
			# two conv one Maxpooling
			nn.Conv2d(32,64,3, padding=1, stride=2),
			nn.BatchNorm2d(64),
			nn.ReLU(),
			nn.Dropout(0.3),
			# 103 d

			nn.Conv2d(64, 64, 3, padding=1, stride=1),
			nn.BatchNorm2d(64),
			nn.ReLU(),
			nn.Dropout(0.3),
			# 103 d

			# maxpooling
			nn.MaxPool2d(kernel_size=(2, 2), stride=2),
			# 52 d

			nn.Conv2d(64, 128, 2, padding=1, stride=2),
			nn.BatchNorm2d(128),
			nn.ReLU(),
			nn.Dropout(0.3),
			# 27 d
			nn.Conv2d(128, 128,2, padding=1, stride=1),
			nn.BatchNorm2d(128),
			nn.ReLU(),
			nn.Dropout(0.3),
			# 27 d

			# maxpooling
			nn.MaxPool2d(kernel_size=(2, 2), stride=2),
			nn.ReLU(),
			# 13 d

			nn.Conv2d(128, 256, 3, padding=1, stride=2),
			nn.BatchNorm2d(256),
			nn.ReLU(),
			nn.Dropout(0.3),

			nn.Conv2d(256, 256, 3, padding=1, stride=1),
			nn.BatchNorm2d(256),
			nn.ReLU(),
			nn.Dropout(0.3),
			nn.MaxPool2d(kernel_size=(2, 2), stride=1, padding=1),
			nn.ReLU(),
                        
			# 8 d
		)
		self.Liner_layer = nn.Sequential(
			nn.Flatten(),
			nn.Dropout(0.3),
			nn.ReLU(),
			nn.Linear(16384, 1)
		)

		self.ca_1 = CA_attention(32,206,206)
		self.ca_2 = CA_attention(256,8,8)

	def forward(self,x):
		out = self.Frist_layer(x)
		out = self.ca_1(out)
		out = self.hidden_layer(out)
		out = self.ca_2(out)
		out = self.Liner_layer(out)
		return out

